package cryptoutil

import (
	"crypto/ecdsa"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"github.com/pkg/errors"
)

func getCertFromPem(idBytes []byte) (*x509.Certificate, error) {
	if idBytes == nil {
		return nil, errors.New("getCertFromPem error: nil idBytes")
	}

	// Decode the pem bytes
	pemCert, _ := pem.Decode(idBytes)
	if pemCert == nil {
		return nil, errors.Errorf("getCertFromPem error: could not decode pem bytes [%v]", idBytes)
	}

	// get a cert
	var cert *x509.Certificate
	cert, err := x509.ParseCertificate(pemCert.Bytes)
	if err != nil {
		return nil, errors.Wrap(err, "getCertFromPem error: failed to parse x509 cert")
	}

	return cert, nil
}

func GetPrivateKeyFromPEM(raw []byte, pwd []byte) (Key, error) {
	k, err := pemToPrivateKey(raw, pwd)
	if err != nil {
		return nil, err
	}
	switch ks := k.(type) {
	case *ecdsa.PrivateKey:
		return &ecdsaPrivateKey{ks}, nil
	default:
		return nil, fmt.Errorf("unknown key type")
	}
}

func pemToPrivateKey(raw []byte, pwd []byte) (interface{}, error) {
	b, _ := pem.Decode(raw)
	if b == nil {
		return nil, fmt.Errorf("")
	}
	if x509.IsEncryptedPEMBlock(b) {
		// TODO
		if len(pwd) == 0 {
			return nil, fmt.Errorf("must have password")
		}
		return encryptoPEMToPrivateKey(raw, pwd)
	}

	derBytes := b.Bytes
	if k, err := x509.ParsePKCS1PrivateKey(derBytes); err == nil {
		return k, nil
	}

	if k, err := x509.ParsePKCS8PrivateKey(derBytes); err == nil {
		return k, nil
	}

	if k, err := x509.ParseECPrivateKey(derBytes); err == nil {
		return k, nil
	}
	return nil, errors.New("Invalid key type. The DER must contain an rsa.PrivateKey or ecdsa.PrivateKey")
}


func encryptoPEMToPrivateKey(raw []byte, pwd []byte) (interface{}, error) {
	if len(raw) == 0 {
		return nil, errors.New("Invalid PEM. It must be different from nil.")
	}
	block, _ := pem.Decode(raw)
	if block == nil {
		return nil, fmt.Errorf("Failed decoding PEM. Block must be different from nil. [% x]", raw)
	}
	if x509.IsEncryptedPEMBlock(block) {
		if len(pwd) == 0 {
			return nil, errors.New("Encrypted Key. Need a password")
		}

		decrypted, err := x509.DecryptPEMBlock(block, pwd)
		if err != nil {
			return nil, fmt.Errorf("Failed PEM decryption [%s]", err)
		}

		key, err := DERToPrivateKey(decrypted)
		if err != nil {
			return nil, err
		}
		return key, err
	}
	cert, err := DERToPrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	return cert, err
}

// DERToPrivateKey unmarshals a der to private key
func DERToPrivateKey(der []byte) (key interface{}, err error) {

	if key, err = x509.ParsePKCS1PrivateKey(der); err == nil {
		return key, nil
	}

	if key, err = x509.ParsePKCS8PrivateKey(der); err == nil {
		switch key.(type) {
		case *rsa.PrivateKey, *ecdsa.PrivateKey:
			return
		default:
			return nil, errors.New("Found unknown private key type in PKCS#8 wrapping")
		}
	}

	if key, err = x509.ParseECPrivateKey(der); err == nil {
		return
	}

	return nil, errors.New("Invalid key type. The DER must contain an rsa.PrivateKey or ecdsa.PrivateKey")
}
